import sys

sys.path.append('..')
from utils.log import Logger
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
import xgboost
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import AdaBoostClassifier
from sklearn.metrics import precision_score, recall_score, f1_score
import matplotlib.pyplot as plt
import joblib
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import StandardScaler
from sklearn.utils import class_weight
import os
import datetime
import numpy as np

os.environ['OMP_NUM_THREADS'] = '1'
from sklearn.cluster import KMeans


# 数据特征与标签关系分析以及展示
def data_plt(data):
    import warnings

    # 忽略 UserWarning
    warnings.filterwarnings("ignore", category=UserWarning)

    plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体字体
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示为方块的问题

    print(data.info())
    print(f'0-4{data.iloc[:11, :4]}')
    print(f"5-8{data.iloc[:11, 4:8]}")
    print(f"9-12{data.iloc[:11, 8:12]}")
    print(f"13-16{data.iloc[:11, 12:16]}")
    print(f"17-20{data.iloc[:11, 16:20]}")
    print(f"21{data.iloc[:11, 20:21]}")

    print(f'出差情况{data.groupby('BusinessTravel')['Attrition'].value_counts()}')
    print(f'部门情况{data.groupby('Department')['Attrition'].value_counts()}')
    print(f'离家距离情况{data.groupby('DistanceFromHome')['Attrition'].value_counts()}')
    print(f'教育程度情况{data.groupby('Education')['Attrition'].value_counts()}')
    print(f'教育领域情况{data.groupby('EducationField')['Attrition'].value_counts()}')
    print(f'环境满意度情况{data.groupby('EnvironmentSatisfaction')['Attrition'].value_counts()}')
    print(f'工作投入情况{data.groupby('JobInvolvement')['Attrition'].value_counts()}')
    print(f'职位等级情况{data.groupby('JobLevel')['Attrition'].value_counts()}')
    print(f'工作岗位情况{data.groupby('JobRole')['Attrition'].value_counts()}')
    print(f'工作满意度情况{data.groupby('JobSatisfaction')['Attrition'].value_counts()}')
    print(f'婚姻状况情况{data.groupby('MaritalStatus')['Attrition'].value_counts()}')
    print(f'月收入情况{data.groupby('MonthlyIncome')['Attrition'].value_counts()}')
    print(f'服务过几家公司情况{data.groupby('NumCompaniesWorked')['Attrition'].value_counts()}')
    print(f'加班情况{data.groupby('OverTime')['Attrition'].value_counts()}')
    print(f'关系满意度情况{data.groupby('RelationshipSatisfaction')['Attrition'].value_counts()}')
    print(f'股权等级情况{data.groupby('StockOptionLevel')['Attrition'].value_counts()}')
    print(f'总工作年度情况{data.groupby('TotalWorkingYears')['Attrition'].value_counts()}')
    print(f'上一次培训次数情况{data.groupby('TrainingTimesLastYear')['Attrition'].value_counts()}')
    print(f'工作生活平衡情况{data.groupby('WorkLifeBalance')['Attrition'].value_counts()}')
    print(f'工作时长情况{data.groupby('YearsAtCompany')['Attrition'].value_counts()}')
    print(f'当前岗位在职时长情况{data.groupby('YearsInCurrentRole')['Attrition'].value_counts()}')
    print(f'上次升职时间情况{data.groupby('YearsSinceLastPromotion')['Attrition'].value_counts()}')
    print(f'和现任经理时长情况{data.groupby('YearsWithCurrManager')['Attrition'].value_counts()}')

    # 定义分组列和对应的描述
    columns = [
        ('BusinessTravel', '出差情况'),
        ('Department', '部门情况'),
        ('DistanceFromHome', '离家距离情况'),
        ('Education', '教育程度情况'),
        ('EducationField', '教育领域情况'),
        ('EnvironmentSatisfaction', '环境满意度情况'),
        ('JobInvolvement', '工作投入情况'),
        ('JobLevel', '职位等级情况'),
        ('JobRole', '工作岗位情况'),
        ('JobSatisfaction', '工作满意度情况'),
        ('MaritalStatus', '婚姻状况情况'),
        ('MonthlyIncome', '月收入情况'),
        ('NumCompaniesWorked', '服务过几家公司情况'),
        ('OverTime', '加班情况'),
        ('RelationshipSatisfaction', '关系满意度情况'),
        ('StockOptionLevel', '股权等级情况'),
        ('TotalWorkingYears', '总工作年度情况'),
        ('TrainingTimesLastYear', '上一次培训次数情况'),
        ('WorkLifeBalance', '工作生活平衡情况'),
        ('YearsAtCompany', '工作时长情况'),
        ('YearsInCurrentRole', '当前岗位在职时长情况'),
        ('YearsSinceLastPromotion', '上次升职时间情况'),
        ('YearsWithCurrManager', '和现任经理时长情况')
    ]

    # 计算子图布局
    nrows = 6
    ncols = 4
    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(20, 30))
    plt.subplots_adjust(hspace=0.5, wspace=0.3)

    # 循环处理每个分组列
    for i, (col, desc) in enumerate(columns):
        result = data.groupby(col)['Attrition'].value_counts()
        reshaped_result = result.unstack(fill_value=0)
        bar_width = 0.35
        index = np.arange(len(reshaped_result.index))

        row = i // ncols
        col_idx = i % ncols
        ax = axes[row, col_idx]

        ax.bar(index - bar_width / 2, reshaped_result[0], bar_width, label='Attrition 0')
        ax.bar(index + bar_width / 2, reshaped_result[1], bar_width, label='Attrition 1')

        ax.set_xlabel(col)
        ax.set_ylabel('数量')
        ax.set_title(f'{desc}Attrition 分布')
        ax.set_xticks(index)
        ax.set_xticklabels(reshaped_result.index, rotation=45)
        ax.legend()

    # 隐藏多余的子图
    for j in range(i + 1, nrows * ncols):
        row = j // ncols
        col_idx = j % ncols
        axes[row, col_idx].axis('off')
    plt.savefig('../diagram/all_grouped_charts.png', dpi=300, bbox_inches='tight')
    plt.show()


# 数据处理&特征删除&特征热编码
def Data_process(data):
    feature = ['Attrition', 'Age', 'BusinessTravel', 'Department', 'DistanceFromHome',
               'Education', 'EducationField', 'EmployeeNumber',
               'EnvironmentSatisfaction', 'Gender', 'JobInvolvement', 'JobLevel',
               'JobRole', 'JobSatisfaction', 'MaritalStatus', 'MonthlyIncome',
               'NumCompaniesWorked', 'Over18', 'OverTime', 'PercentSalaryHike',
               'PerformanceRating', 'RelationshipSatisfaction', 'StandardHours',
               'StockOptionLevel', 'TotalWorkingYears', 'TrainingTimesLastYear',
               'WorkLifeBalance', 'YearsAtCompany', 'YearsInCurrentRole',
               'YearsSinceLastPromotion', 'YearsWithCurrManager']
    ret_data = data.copy(deep=True)
    ret_data.drop(
        ['Over18', 'NumCompaniesWorked', 'PercentSalaryHike', 'TrainingTimesLastYear', 'EmployeeNumber', 'Education',
         'EducationField'],
        axis=1, inplace=True)  # 删除掉不重要的特征列和员工号
    OneHot_lis = ['BusinessTravel', 'Department', 'Gender', 'JobLevel',
                  'JobSatisfaction', 'MaritalStatus', 'EnvironmentSatisfaction', 'StockOptionLevel',
                  'RelationshipSatisfaction', 'WorkLifeBalance', 'JobRole']
    ret_data['OverTime'] = data['OverTime'].apply(lambda x: 0 if x == 'No' else 1)
    enc = OneHotEncoder(sparse_output=False)

    encoded_data = enc.fit_transform(data[feature][OneHot_lis])
    # print(enc_BusinessTravel.get_feature_names_out())
    data_OneHot = pd.DataFrame(encoded_data, columns=enc.get_feature_names_out())
    ret_data = pd.concat([ret_data, data_OneHot], axis=1)
    joblib.dump(enc, '../model/OneHot_enc.pkl')
    ret_data.drop(OneHot_lis, axis=1, inplace=True)
    print(ret_data)
    print(ret_data.info())
    print(ret_data.iloc[:, 1:])

    return ret_data


# 模型训练以及模型评估
def model_train(data, logger):
    x = data.iloc[:, 1:]
    y = data['Attrition']

    # 数据标准化
    scaler = StandardScaler()
    x = scaler.fit_transform(x)
    logger.info("===================开始LogisticRegression模型训练===================")
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=24)
    es_log = LogisticRegression(max_iter=1000, random_state=24)
    es_log.fit(x_train, y_train)
    y_pre_log = es_log.predict(x_test)

    print(f'精确率：{precision_score(y_test, y_pre_log, pos_label=0)}')

    print(f'召回率：{recall_score(y_test, y_pre_log, pos_label=0)}')

    print(f'f1：{f1_score(y_test, y_pre_log, pos_label=0)}')
    y_score_log = es_log.predict_proba(x_test)[:, 1]
    print(f'逻辑回归预测AUC：{roc_auc_score(y_test, y_score_log)}')
    logger.info("=========================模型训练完成=============================")
    logger.info(f"LogisticRegression模型在训练集上的AUC：{roc_auc_score(y_test, y_score_log)}")

    # # es_xgb = xgboost()
    # logger.info("=========开始单一决策树模型训练===================")
    # es_tree = DecisionTreeClassifier()
    # es_tree.fit(x_train, y_train)
    # y_pre_tree = es_tree.predict(x_test)
    # print(f"单一决策树正确率{accuracy_score(y_test, y_pre_tree)}")
    # y_score_tree = es_log.predict_proba(x_test)[:, 1]
    # print(f'单一决策树预测AUC：{roc_auc_score(y_test, y_score_tree)}')
    # logger.info("=========================模型训练完成=============================")
    # logger.info(f"单一决策树模型在训练集上的AUC：{roc_auc_score(y_test, y_score_tree)}")
    #
    # logger.info("=========开始AadBoost集成学习决策树模型训练===================")
    # es_ada_tree = AdaBoostClassifier(estimator=es_tree, n_estimators=100, learning_rate=1.5)
    # es_ada_tree.fit(x_train, y_train)
    # y_pre_ada_tree = es_ada_tree.predict(x_test)
    # print(f"AadBoost集成学习决策树预测正确率{accuracy_score(y_test, y_pre_ada_tree)}")
    # y_score_ada_tree = es_log.predict_proba(x_test)[:, 1]
    # print(f'AadBoost集成学习决策树预测AUC：{roc_auc_score(y_test, y_score_ada_tree)}')
    # logger.info("=========================模型训练完成=============================")
    # logger.info(f"AadBoost集成学习决策树在训练集上的AUC：{roc_auc_score(y_test, y_score_ada_tree)}")
    #
    # logger.info("=========开始AadBoost集成学习逻辑回归模型训练===================")
    # es_ada_log = AdaBoostClassifier(estimator=es_log, n_estimators=100, learning_rate=1.5)
    # es_ada_log.fit(x_train, y_train)
    # y_pre_ada_log = es_ada_log.predict(x_test)
    # print(f"AadBoost集成学习逻辑回归预测正确率{accuracy_score(y_test, y_pre_ada_log)}")
    # y_score_ada_log = es_ada_log.predict_proba(x_test)[:, 1]
    # print(f'AadBoost集成学习逻辑回归预测AUC：{roc_auc_score(y_test, y_score_ada_log)}')
    # logger.info("=========================模型训练完成=============================")
    # logger.info(f"AadBoost集成学习逻辑回归模型在训练集上的AUC：{roc_auc_score(y_test, y_score_ada_log)}")

    # logger.info("=========开始xgboost模型训练===================")
    # es_xgb = xgboost.XGBClassifier(
    #     n_estimators=100,  # 树的数量
    #     learning_rate=0.1,  # 学习率
    #     max_depth=10,  # 树的最大深度
    #     random_state=24,  # 随机种子
    #     use_label_encoder=False,  # 新版本中不需要 label encoder
    #     eval_metric='logloss')
    # class_weight.compute_sample_weight('balanced', y_train)
    # es_xgb.fit(x_train, y_train)
    # y_pre_ada_log = es_xgb.predict(x_test)
    # print(f"xgboost预测正确率{accuracy_score(y_test, y_pre_ada_log)}")
    # y_score_ada_log = es_xgb.predict_proba(x_test)[:, 1]
    # print(f'xgboost预测AUC：{roc_auc_score(y_test, y_score_ada_log)}')
    # logger.info("=========================模型训练完成=============================")
    # logger.info(f"xgboost模型在训练集上的AUC：{roc_auc_score(y_test, y_score_ada_log)}")

    joblib.dump(es_log, '../model/logist.pkl')
    # joblib.dump(es_ada_tree, '../model/ada_tree.pkl')
    # joblib.dump(es_tree, '../model/tree.pkl')
    # joblib.dump(es_ada_log, '../model/ada_log.pkl')
    # joblib.dump(es_xgb, '../model/xgb.pkl')
    logger.info("=========================模型已保存=============================")


if __name__ == '__main__':
    data_pd = pd.read_csv('../data/train.csv')
    data_copy = data_pd.copy(deep=True)
    logger = Logger(root_path='../log', log_name="train_" + datetime.datetime.now().strftime('%Y%m%d%H%M%S'),
                    level='debug').get_logger()
    # data_plt(data_copy)
    data_train_ana = Data_process(data_copy)
    model_train(data_train_ana, logger)
